#!/bin/env python
from netCDF4 import Dataset
import numpy as np
import math

#
# DGPI computation using a renalayis dataset
#

#--Directory Setting
input_reanalysis_dir = "./input_data/reanalysis" # input directory for reanalysis data
input_besttrack_dir = "./input_data/besttrack" # input directory for tropical cyclone tracks
output_dir = "./output_data" # output directory
output_file1 = "%s/dgpi.nc" % (output_dir) # output: monthly DGPI data
output_file2 = "%s/dgpi_climmonthly.nc" % (output_dir) # output: climatological monthly DGPI data

#--Parameters Setting
##--DGPI coefficients
coef={
  'dudy'    :2.3, # original: 2.0
  'wshear'  :-1.7, # original: -1.0
  'avor'    :2.4,  # original: 2.0
  'omega'   :3.4,  # original: 3.0
  'const'   :-11.8, # original: -12.0
}
undef = -9.99E33 # undefined value

#--Read inputs
##--dudy at 500hPa
f1 = Dataset("%s/msu500.nc" % (input_reanalysis_dir), "r")
dudy = f1.variables['msu500']

##--omega at 500hPa
f2 = Dataset("%s/omega500.nc" % (input_reanalysis_dir), "r")
omega = f2.variables['omega500']

##--absolute vorticity at 850 hPa
f3 = Dataset("%s/avor850a.nc" % (input_reanalysis_dir), "r")
avor = f3.variables['avor850a']

##--vertical wind shear
f4 = Dataset("%s/wshear.nc" % (input_reanalysis_dir), "r")
wshear = f4.variables['wshear']

##--relative SST anomaly
f5 = Dataset("%s/sstanm.nc" % (input_reanalysis_dir), "r")
ssta = f5.variables['sstanm']

##--grids, longitudes, and latitudes
tmax,jmax,imax = np.shape(ssta[:])

lon_in = f5.variables['lon']
lat_in = f5.variables['lat']
times_in = f5.variables['time']

lons,lats = np.meshgrid(lon_in[:],lat_in[:])
lats_tile = np.tile(lats,(tmax,1,1))

#-----Compute DGPI---------------------------
dudy_term = (5.5 - 100000*dudy[:])**(coef['dudy'])
wshear_term = (2.0 + 0.1*wshear[:])**(coef['wshear'])
omega_term = (5.0 - 20*omega[:])**(coef['omega'])
avor_term = (5.5 + 100000*avor[:])**(coef['avor'])
const_term = math.exp(coef['const'])
dgpi = -1.0 + dudy_term * wshear_term * omega_term * avor_term * const_term

##-----Post Process---------------------------
##--Fill zero value where relative SST anomaly is negative
neg_ssta_index = np.where(ssta[:]<0)
dgpi[neg_ssta_index] = 0.0
dudy_term[neg_ssta_index] = 0.0
wshear_term[neg_ssta_index] = 0.0
omega_term[neg_ssta_index] = 0.0
avor_term[neg_ssta_index] = 0.0
 
##--Fill zero value where latitude is smaller than 2 degree
neg_lat_index = np.where((lats_tile<2)&(lats_tile>-2))
dgpi[neg_lat_index] = 0.0
dudy_term[neg_lat_index] = 0.0
wshear_term[neg_lat_index] = 0.0
omega_term[neg_lat_index] = 0.0
avor_term[neg_lat_index] = 0.0
 
##--Replace negative DGPI value with zero
#neg_index = np.where(dgpi<0)
#dgpi[neg_index] = 0.0
#dudy_term[neg_index] = 0.0
#wshear_term[neg_index] = 0.0
#omega_term[neg_index] = 0.0
#avor_term[neg_index] = 0.0
 
##--Mask over land surface
dgpi = np.ma.masked_where(ssta[:].mask==True, dgpi)
dudy_term = np.ma.masked_where(ssta[:].mask==True, dudy_term)
wshear_term = np.ma.masked_where(ssta[:].mask==True, wshear_term)
omega_term = np.ma.masked_where(ssta[:].mask==True, omega_term)
avor_term = np.ma.masked_where(ssta[:].mask==True, avor_term)

#print dgpi,np.shape(dgpi),type(dgpi)

#-----Output---------------------------------
#-----Monthly
f6 = Dataset(output_file1, "w", format="NETCDF3_CLASSIC")

#dimmensions
time = f6.createDimension('time', None)
lat = f6.createDimension('lat',jmax)
lon = f6.createDimension('lon',imax)

#variables
times = f6.createVariable('time','f8',('time',))
latitudes = f6.createVariable('lat','f4', ('lat',))
longitudes = f6.createVariable('lon','f4', ('lon',))

dgpi_out = f6.createVariable('dgpi','f4',('time','lat','lon'),fill_value=undef)
dudy_out = f6.createVariable('dudy','f4',('time','lat','lon'),fill_value=undef)
wshear_out = f6.createVariable('wshear','f4',('time','lat','lon'),fill_value=undef)
omega_out = f6.createVariable('omega','f4',('time','lat','lon'),fill_value=undef)
avor_out = f6.createVariable('avor','f4',('time','lat','lon'),fill_value=undef)
    
#attributtes
f6.description = 'DGPI'
latitudes.units = lat_in.units
latitudes.axis= 'Y'
longitudes.units = lon_in.units
longitudes.axis= 'X'
times.units = times_in.units
times.calendar = times_in.calendar
times.cargtesian_axis = "T"

dgpi_out.long_name = "Monthly DGPI"
dudy_out.long_name = "DUDY Term"
wshear_out.long_name = "Vertical Wind Shear Term"
omega_out.long_name = "Omega Term"
avor_out.long_name = "Absolute Vorticity Term"

latitudes[:] = lat_in[:]
longitudes[:] = lon_in[:]
times[:] = times_in[:]

dgpi_out[:] = dgpi[:]
dudy_out[:] = dudy_term[:]
wshear_out[:] = wshear_term[:]
omega_out[:] = omega_term[:]
avor_out[:] = avor_term[:]

#-----Climatological Monthly
f7 = Dataset(output_file2, "w", format="NETCDF3_CLASSIC")

#dimmensions
time = f7.createDimension('time', 12)
lat = f7.createDimension('lat',jmax)
lon = f7.createDimension('lon',imax)

#variables
times = f7.createVariable('time','f8',('time',))
latitudes = f7.createVariable('lat','f4', ('lat',))
longitudes = f7.createVariable('lon','f4', ('lon',))

dgpi_out = f7.createVariable('dgpi','f4',('time','lat','lon'),fill_value=undef)
dudy_out = f7.createVariable('dudy','f4',('time','lat','lon'),fill_value=undef)
wshear_out = f7.createVariable('wshear','f4',('time','lat','lon'),fill_value=undef)
omega_out = f7.createVariable('omega','f4',('time','lat','lon'),fill_value=undef)
avor_out = f7.createVariable('avor','f4',('time','lat','lon'),fill_value=undef)
    
#attributtes
f7.description = 'DGPI'
latitudes.units = lat_in.units
latitudes.axis= 'Y'
longitudes.units = lon_in.units
longitudes.axis= 'X'
times.units = times_in.units
times.calendar = times_in.calendar
times.cargtesian_axis = "T"

dgpi_out.long_name = "Climatological Monthly DGPI"
dudy_out.long_name = "DUDY Term"
wshear_out.long_name = "Vertical Wind Shear Term"
omega_out.long_name = "Omega Term"
avor_out.long_name = "Absolute Vorticity Term"

latitudes[:] = lat_in[:]
longitudes[:] = lon_in[:]
times[:] = times_in[0:12]

dgpi_out[:] = dgpi[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)
dudy_out[:] = dudy_term[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)
wshear_out[:] = wshear_term[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)
omega_out[:] = omega_term[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)
avor_out[:] = avor_term[:].reshape(tmax/12,12,jmax,imax).mean(axis=0)

#-----Finalize
f1.close()
f2.close()
f3.close()
f4.close()
f5.close()
f6.close()
f7.close()
